#include <vector>
#include <algorithm>
#include "informedRRT.h"
#include "KDTree.hpp"
#include "KDNode.hpp"

#define INF 1000000.0

using std::vector;

namespace global_planner{

/**
 * @brief 计算两个节点的欧氏距离
 */
double informedRRT::GetEuclideanDistance(const KDNodePtr& p1, const KDNodePtr& p2){
    return sqrt(pow(p1->point[X] - p2->point[X], 2) + pow(p1->point[Y] - p2->point[Y], 2));
}
/**
 * @brief 检测两个节点间连线是否有障碍物
 */
bool informedRRT::isCollision(unsigned char* costs, const KDNodePtr& p1, const KDNodePtr& p2, double maxDist){
    double theta = atan2(p2->point[Y] - p1->point[Y], p2->point[X] - p1->point[X]); // 倾角
    double dist = GetEuclideanDistance(p1, p2); // 两点距离
    if(dist > maxDist)
    {
        return true;
    }
    double step = 1; // 步长
    int n_step = (int)(dist / step); // 步数
    double px = p1->point[X] + step * cos(theta);
    double py = p1->point[Y] + step * sin(theta);

    auto n = std::make_shared<KDNode>((int)round(px), (int)round(py));
    for (int i = 0; i < n_step; i++)
    {
        n->point[X] = (int)px;
        n->point[Y] = (int)py;
        if(costs[getNodeIndex(n)] != costmap_2d::FREE_SPACE)
        {
            return true;  // 有碰撞
        }

        px += step * cos(theta);
        py += step * sin(theta);
    }
    return false;
}
/**
 * @brief 检测两个节点间连线是否有障碍物
 */
bool informedRRT::isCollision(unsigned char* costs, const KDNodePtr& p1, const KDNodePtr& p2){
    double theta = atan2(p2->point[Y] - p1->point[Y], p2->point[X] - p1->point[X]); // 倾角
    double dist = GetEuclideanDistance(p1, p2); // 两点距离

    double step = 1; // 步长
    int n_step = (int)(dist / step); // 步数
    double px = p1->point[X] + step * cos(theta);
    double py = p1->point[Y] + step * sin(theta);

    auto n = std::make_shared<KDNode>((int)round(px), (int)round(py));
    for (int i = 0; i < n_step; i++)
    {
        n->point[X] = (int)px;
        n->point[Y] = (int)py;
        if(costs[getNodeIndex(n)] != costmap_2d::FREE_SPACE)
        {
            return true;  // 有碰撞
        }

        px += step * cos(theta);
        py += step * sin(theta);
    }
    return false;
}
/**
 * @brief 在指定范围内生成一个随机整数
 */
int random(int min, int max){
    return (int)(std::rand() / (double)RAND_MAX * (max - min) + min);
}

/**
 * @brief 将在椭圆内采样得到的点旋转到世界坐标系
 */
KDNodePtr informedRRT::RotationToWorldFrame(const KDNodePtr& sample, const Hyperellipsoid& hyperellipsoid)
{
    auto sample_new = std::make_shared<KDNode>( (int)(hyperellipsoid.rotationMatrix[0][0] * sample->point[X] + hyperellipsoid.rotationMatrix[0][1] * sample->point[Y] + hyperellipsoid.centre->point[X]),
                                                (int)(hyperellipsoid.rotationMatrix[1][0] * sample->point[X] + hyperellipsoid.rotationMatrix[1][1] * sample->point[Y] + hyperellipsoid.centre->point[Y]));
    sample_new->point[X] = sample_new->point[X] < 0 ? 0 : sample_new->point[X];
    sample_new->point[X] = sample_new->point[X] >= nx_ ? (nx_ - 1) : sample_new->point[X];
    sample_new->point[Y] = sample_new->point[Y] < 0 ? 0 : sample_new->point[Y];
    sample_new->point[Y] = sample_new->point[Y] >= ny_ ? (ny_ - 1) : sample_new->point[Y];

    return sample_new;
}

/**
 * @brief 在椭圆范围内采样
 * @param  costs            代价地图
 * @param  hyperellipsoid   椭圆参数
 */
KDNodePtr informedRRT::GetSample(unsigned char* costs, const Hyperellipsoid& hyperellipsoid, double c_max)
{
    auto sample = std::make_shared<KDNode>(random(0, nx_), random(0, ny_));
    if (c_max == INF)
    {
        while (costs[getNodeIndex(sample)] != costmap_2d::FREE_SPACE)
        {
            sample = std::make_shared<KDNode>(random(0, nx_), random(0, ny_));
        }
        return sample;
    }
    double x_max = c_max / 2; // 长轴

    double c = hyperellipsoid.c_min / 2.0; // 焦点坐标
    double y_max = sqrt(x_max * x_max + c * c); // 短轴
    
    auto focus_neg = std::make_shared<KDNode>(-c, 0); // 焦点
    auto focus_pos = std::make_shared<KDNode>( c, 0); // 焦点
    do
    {
        auto sample0 = std::make_shared<KDNode>(random((int)-x_max, (int)x_max), random((int)-y_max, (int)y_max));
        // 判断采样点是否在椭圆内
        // 如果采样点到两焦点的距离之和小于2a，则采样点在椭圆内
        while (GetEuclideanDistance(sample0, focus_neg) + GetEuclideanDistance(sample0, focus_pos) > 2 * x_max)
        {
            sample0 = std::make_shared<KDNode>(random((int)-x_max, (int)x_max), random((int)-y_max, (int)y_max));
        }

        sample = RotationToWorldFrame(sample0, hyperellipsoid);
    } while (costs[getNodeIndex(sample)] != costmap_2d::FREE_SPACE);
    return sample;
}
/**
 * @brief 搜索函数
 * @param  costs   代价地图
 * @param  start   起点
 * @param  goal    终点
 * @param  path    搜索到的路径
 * @return true 
 * @return false 
 */
bool informedRRT::search(   unsigned char* costs,
                            const KDNodePtr& start,
                            const KDNodePtr& goal,
                            vector<KDNodePtr>& path){

    ROS_INFO("start point: %.1f, %.1f", start->point[X], start->point[Y]);
    ROS_INFO("goal  point: %.1f, %.1f", goal->point[X], goal->point[Y]);
    const int ITER = 2000;
    const double MAX_DISTANCE = 10;
    const double OptimizationR = 30;
    Hyperellipsoid hyperellipsoid(start, goal);

    KDTree kdtree;
    kdtree.Add(start);
    double c_best = INF;
    int cnt = 0;
    for (int iter = 0; iter < ITER; iter++)
    {
        // 保存原始采样点
        auto sample = GetSample(costs, hyperellipsoid, c_best); 

        // 最近邻搜索
        sample->front = kdtree.kdnodes[0];
        double cost = GetEuclideanDistance(sample, kdtree.kdnodes[0]);
        for(int i = 1; i < kdtree.kdnodes.size(); i++)
        {
            auto n = kdtree.kdnodes[i];
            double new_cost = GetEuclideanDistance(sample, n); // 计算代价
            if (new_cost < cost) // 如果新路径代价更小
            {
                sample->front = n;
                cost = new_cost;
            }
        } 
        double distance_near = GetEuclideanDistance(sample->front, sample);
        if (distance_near > MAX_DISTANCE)
        {
            // 如果原始采样点与当前节点距离大于最大距离
            // 则将两节点连线上距离当前节点MAX_DISTANCE远的节点作为新采样点
            // 也就是当前节点朝原始采样点移动MAX_DISTANCE后得到的节点
            double theta = atan2(sample->point[Y] - sample->front->point[Y], sample->point[X] - sample->front->point[X]); // 倾角
            sample->point[X] = sample->front->point[X] + (int)(MAX_DISTANCE * cos(theta));
            sample->point[Y] = sample->front->point[Y] + (int)(MAX_DISTANCE * sin(theta));
        }

        if (!isCollision(costs, sample->front, sample)) // 如果采样点和已知样点图间存在一条无碰撞路径
        {
            kdtree.Add(sample);
            sample->cost = sample->front->cost + GetEuclideanDistance(sample, sample->front); // 计算代价

            // 重连 Rewire
            double R = OptimizationR;
            vector<KDNodePtr> nearests;
            for(int i = 0; i < kdtree.kdnodes.size(); i++)
            {
                auto n = kdtree.kdnodes[i];
                double cost = n->cost + GetEuclideanDistance(sample, n); // 计算代价
                if (cost < sample->cost) // 如果新路径代价更小
                {
                    if (!isCollision(costs, sample, n, R)) // 碰撞检测通过
                    {
                        sample->cost = cost; // 更新代价
                        sample->front = n; // 更新路径
                    }
                }
                else
                {
                    cost = sample->cost + GetEuclideanDistance(sample, n); // 计算代价
                    if (cost < n->cost) // 如果新路径代价更小
                    {
                        if (!isCollision(costs, sample, n, R)) // 碰撞检测通过
                        {
                            n->cost = cost; // 更新代价
                            n->front = sample; // 更新路径
                        }
                    }
                }

            }

            // 到达终点附近，回溯得到路径
            if (!isCollision(costs, sample, goal, MAX_DISTANCE))
            {
                double c_new = sample->cost + GetEuclideanDistance(goal, sample);

                if (c_new < c_best) // 新路径的长度更小
                {
                    // 回溯路径
                    path.clear();
                    path.push_back(goal);
                    while (sample->front != nullptr)
                    {
                        path.push_back(sample);
                        sample = sample->front;
                    }
                    path.push_back(start);
                    std::reverse(path.begin(), path.end());
                    c_best = c_new;

                    // 剔除超椭圆外的样点
                    vector<KDNodePtr> kdnodes;
                    for (int i = 0; i < kdtree.kdnodes.size(); i++)
                    {
                        auto n = kdtree.kdnodes[i];
                        if (GetEuclideanDistance(n, start)
                            + GetEuclideanDistance(n, goal) < c_best)
                        {
                            kdnodes.push_back(n);
                        }
                    }
                    
                    kdtree.kdnodes = kdnodes;
                }
            }
        }
    }

    return true;

}
};